use chrono::Local;
use colored::*;
use image::{DynamicImage, GenericImageView, ImageBuffer, Rgba, RgbaImage};
use imageproc::drawing::{draw_text_mut, text_size};
use rusttype::{Font, Scale};
use std::fs;
use std::io::{self, Write};
use std::path::{Path, PathBuf};
use walkdir::WalkDir;

// 水印位置枚举
#[derive(Debug, Clone, Copy)]
enum WatermarkPosition {
    TopLeft,
    TopRight,
    BottomLeft,
    BottomRight,
    Center,
}

impl WatermarkPosition {
    fn from_number(n: u32) -> Self {
        match n {
            1 => WatermarkPosition::TopLeft,
            2 => WatermarkPosition::TopRight,
            3 => WatermarkPosition::BottomLeft,
            4 => WatermarkPosition::BottomRight,
            5 => WatermarkPosition::Center,
            _ => WatermarkPosition::BottomRight,
        }
    }

    fn to_string(&self) -> &str {
        match self {
            WatermarkPosition::TopLeft => "左上角",
            WatermarkPosition::TopRight => "右上角",
            WatermarkPosition::BottomLeft => "左下角",
            WatermarkPosition::BottomRight => "右下角",
            WatermarkPosition::Center => "居中",
        }
    }
}

// 水印配置
#[derive(Debug, Clone)]
struct WatermarkConfig {
    text: String,
    font_size: f32,
    opacity: f32,
    position: WatermarkPosition,
    color: Rgba<u8>,
    padding: u32,
}

impl Default for WatermarkConfig {
    fn default() -> Self {
        WatermarkConfig {
            text: "全栈若城".to_string(),
            font_size: 48.0,
            opacity: 0.5,
            position: WatermarkPosition::BottomRight,
            color: Rgba([0, 0, 0, 0]),
            padding: 20,
        }
    }
}

// 图片处理器
struct ImageProcessor {
    config: WatermarkConfig,
    font_data: Vec<u8>,
}

impl ImageProcessor {
    fn new(config: WatermarkConfig) -> Self {
        // 使用阿里巴巴普惠体字体
        let font_data = include_bytes!("../assets/阿里巴巴普惠体.ttf").to_vec();

        ImageProcessor { config, font_data }
    }

    // 添加水印
    fn add_watermark(&self, img: &DynamicImage) -> RgbaImage {
        let mut img_rgba = img.to_rgba8();

        // 加载字体
        let font = match Font::try_from_vec(self.font_data.clone()) {
            Some(f) => f,
            None => {
                println!("{} 字体加载失败，使用默认处理", "[警告]".yellow());
                return img_rgba;
            }
        };

        let scale = Scale::uniform(self.config.font_size);

        // 计算文本尺寸
        let (text_width, text_height) = text_size(scale, &font, &self.config.text);

        // 计算水印位置
        let (width, height) = img.dimensions();
        let (x, y) = self.calculate_position(width as i32, height as i32, text_width, text_height);

        // 应用透明度
        let mut watermark_color = self.config.color;
        watermark_color[3] = (255.0 * self.config.opacity) as u8;

        // 绘制文本
        draw_text_mut(
            &mut img_rgba,
            watermark_color,
            x,
            y,
            scale,
            &font,
            &self.config.text,
        );

        img_rgba
    }

    // 计算水印位置
    fn calculate_position(
        &self,
        img_width: i32,
        img_height: i32,
        text_width: i32,
        text_height: i32,
    ) -> (i32, i32) {
        let padding = self.config.padding as i32;

        match self.config.position {
            WatermarkPosition::TopLeft => (padding, padding),
            WatermarkPosition::TopRight => (img_width - text_width - padding, padding),
            WatermarkPosition::BottomLeft => (padding, img_height - text_height - padding),
            WatermarkPosition::BottomRight => (
                img_width - text_width - padding,
                img_height - text_height - padding,
            ),
            WatermarkPosition::Center => {
                ((img_width - text_width) / 2, (img_height - text_height) / 2)
            }
        }
    }

    // 批量处理图片
    fn process_batch(&self, input_paths: &[PathBuf], output_dir: &str) -> io::Result<Vec<String>> {
        fs::create_dir_all(output_dir)?;
        let mut processed = Vec::new();

        for (idx, input_path) in input_paths.iter().enumerate() {
            println!(
                "\n{} 处理 [{}/{}]: {}",
                ">>".cyan(),
                idx + 1,
                input_paths.len(),
                input_path.display().to_string().bright_black()
            );

            match self.process_single(input_path, output_dir) {
                Ok(output_path) => {
                    println!("{} 完成: {}", "[成功]".green(), output_path.bright_white());
                    processed.push(output_path);
                }
                Err(e) => {
                    println!("{} 失败: {}", "[失败]".red().bold(), e);
                }
            }
        }

        Ok(processed)
    }

    // 处理单张图片
    fn process_single(&self, input_path: &Path, output_dir: &str) -> io::Result<String> {
        // 读取图片
        let img =
            image::open(input_path).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;

        let (width, height) = img.dimensions();
        println!(
            "  {} 尺寸: {}x{} 像素",
            "[信息]".blue(),
            width.to_string().bright_white(),
            height.to_string().bright_white()
        );

        // 添加水印
        let watermarked = self.add_watermark(&img);

        // 生成输出文件名
        let filename = input_path.file_stem().unwrap().to_string_lossy();
        let extension = input_path.extension().unwrap_or_default().to_string_lossy();
        let timestamp = Local::now().format("%Y%m%d_%H%M%S");
        let output_filename = format!("{}_watermark_{}.{}", filename, timestamp, extension);
        let output_path = Path::new(output_dir).join(&output_filename);

        // 保存图片
        watermarked
            .save(&output_path)
            .map_err(|e| io::Error::new(io::ErrorKind::Other, e))?;

        Ok(output_path.display().to_string())
    }
}

// 扫描目录获取图片文件
fn scan_images(dir: &str) -> Vec<PathBuf> {
    let image_extensions = ["jpg", "jpeg", "png", "bmp", "gif", "webp"];

    WalkDir::new(dir)
        .into_iter()
        .filter_map(|e| e.ok())
        .filter(|e| e.path().is_file())
        .filter(|e| {
            e.path()
                .extension()
                .and_then(|ext| ext.to_str())
                .map(|ext| image_extensions.contains(&ext.to_lowercase().as_str()))
                .unwrap_or(false)
        })
        .map(|e| e.path().to_path_buf())
        .collect()
}

fn print_header() {
    println!("\n{}", "═".repeat(60).bright_cyan());
    println!(
        "{}",
        "    Rust 图片水印工具 v1.0    "
            .bright_white()
            .bold()
            .on_blue()
    );
    println!("{}\n", "═".repeat(60).bright_cyan());
}

fn print_menu() {
    println!("{}", "\n主菜单".bright_yellow().bold());
    println!("{}", "─".repeat(40).bright_black());
    println!("  1. 快速添加水印（单张图片）");
    println!("  2. 批量添加水印（整个目录）");
    println!("  3. 自定义水印配置");
    println!("  4. 查看当前配置");
    println!("  5. 创建示例图片");
    println!("  0. 退出程序");
    println!("{}", "─".repeat(40).bright_black());
    print!("\n请选择操作 (0-5): ");
    io::stdout().flush().unwrap();
}

fn get_input() -> String {
    let mut input = String::new();
    io::stdin().read_line(&mut input).unwrap();
    input.trim().to_string()
}

fn configure_watermark(config: &mut WatermarkConfig) {
    println!("\n{}", " 自定义水印配置".bright_cyan().bold());

    print!("\n水印文字 (当前: {}): ", config.text.bright_white());
    io::stdout().flush().unwrap();
    let text = get_input();
    if !text.is_empty() {
        config.text = text;
    }

    print!(
        "字体大小 (当前: {}, 建议 20-100): ",
        config.font_size.to_string().bright_white()
    );
    io::stdout().flush().unwrap();
    if let Ok(size) = get_input().parse::<f32>() {
        config.font_size = size.clamp(10.0, 200.0);
    }

    print!(
        "透明度 (当前: {}, 范围 0.0-1.0): ",
        config.opacity.to_string().bright_white()
    );
    io::stdout().flush().unwrap();
    if let Ok(opacity) = get_input().parse::<f32>() {
        config.opacity = opacity.clamp(0.0, 1.0);
    }

    println!("\n水印位置:");
    println!("  1. 左上角");
    println!("  2. 右上角");
    println!("  3. 左下角");
    println!("  4. 右下角");
    println!("  5. 居中");
    print!(
        "选择位置 (当前: {}): ",
        config.position.to_string().bright_white()
    );
    io::stdout().flush().unwrap();
    if let Ok(pos) = get_input().parse::<u32>() {
        config.position = WatermarkPosition::from_number(pos);
    }

    println!("\n水印颜色:");
    println!("  1. 白色 (默认)");
    println!("  2. 黑色");
    println!("  3. 红色");
    println!("  4. 蓝色");
    print!("选择颜色: ");
    io::stdout().flush().unwrap();
    if let Ok(color_choice) = get_input().parse::<u32>() {
        config.color = match color_choice {
            1 => Rgba([255, 255, 255, 255]),
            2 => Rgba([0, 0, 0, 255]),
            3 => Rgba([255, 0, 0, 255]),
            4 => Rgba([0, 100, 255, 255]),
            _ => Rgba([255, 255, 255, 255]),
        };
    }

    println!("\n{} 配置更新完成！", "[成功]".green().bold());
}

fn create_sample_image() -> io::Result<()> {
    fs::create_dir_all("images")?;

    // 创建一个彩色渐变图片作为示例
    let width = 800;
    let height = 600;

    let mut img: RgbaImage = ImageBuffer::new(width, height);

    for y in 0..height {
        for x in 0..width {
            let r = (x as f32 / width as f32 * 255.0) as u8;
            let g = (y as f32 / height as f32 * 255.0) as u8;
            let b = ((x + y) as f32 / (width + height) as f32 * 255.0) as u8;

            img.put_pixel(x, y, Rgba([r, g, b, 255]));
        }
    }

    let output_path = "images/sample.png";
    img.save(output_path)
        .map_err(|e| io::Error::new(io::ErrorKind::Other, e))?;

    println!(
        "\n{} 示例图片已创建: {}",
        "[成功]".green().bold(),
        output_path.bright_cyan()
    );
    println!("  {} 尺寸: {}x{} 像素", "[信息]".blue(), width, height);

    Ok(())
}

fn display_config(config: &WatermarkConfig) {
    println!("\n{}", " 当前水印配置".bright_cyan().bold().underline());
    println!("{}", "─".repeat(40).bright_black());
    println!("  水印文字: {}", config.text.bright_white().bold());
    println!(
        "  字体大小: {}",
        config.font_size.to_string().bright_white()
    );
    println!(
        "  透明度: {} ({}%)",
        config.opacity.to_string().bright_white(),
        (config.opacity * 100.0).to_string().yellow()
    );
    println!("  位置: {}", config.position.to_string().bright_cyan());
    let color_name = match config.color {
        Rgba([255, 255, 255, _]) => "白色",
        Rgba([0, 0, 0, _]) => "黑色",
        Rgba([255, 0, 0, _]) => "红色",
        Rgba([0, 100, 255, _]) => "蓝色",
        _ => "自定义",
    };
    println!("  颜色: {}", color_name.bright_white());
    println!("  边距: {} 像素", config.padding.to_string().bright_white());
    println!("{}\n", "─".repeat(40).bright_black());
}

fn demo_mode() {
    println!("\n{}", "[演示模式]".bright_magenta().bold());
    println!("\n{}", "快速开始指南：".bright_yellow());
    println!("  [1] 选择 {} 创建示例图片", "「5」".bright_green());
    println!("  [2] 选择 {} 自定义水印配置", "「3」".bright_green());
    println!("  [3] 选择 {} 为示例图片添加水印", "「1」".bright_green());
    println!("  [4] 在 {} 目录查看处理结果", "output/".bright_cyan());
    println!();
}

fn main() {
    print_header();
    demo_mode();

    let mut config = WatermarkConfig::default();

    loop {
        print_menu();
        let choice = get_input();

        match choice.as_str() {
            "1" => {
                // 单张图片处理
                print!("\n请输入图片路径: ");
                io::stdout().flush().unwrap();
                let input_path = get_input();

                if input_path.is_empty() {
                    println!("{} 路径不能为空", "[错误]".red().bold());
                    continue;
                }

                let path = Path::new(&input_path);
                if !path.exists() {
                    println!("{} 文件不存在: {}", "[错误]".red().bold(), input_path);
                    continue;
                }

                print!("输出目录 (默认: output): ");
                io::stdout().flush().unwrap();
                let output_dir = get_input();
                let output_dir = if output_dir.is_empty() {
                    "output"
                } else {
                    &output_dir
                };

                let processor = ImageProcessor::new(config.clone());
                println!("\n{}", "[开始] 开始处理...".bright_cyan().bold());

                match processor.process_single(path, output_dir) {
                    Ok(output_path) => {
                        println!("\n{} 处理成功！", "[完成]".bright_green().bold());
                        println!("  输出文件: {}", output_path.bright_cyan());
                    }
                    Err(e) => {
                        println!("{} 处理失败: {}", "[失败]".red().bold(), e);
                    }
                }
            }
            "2" => {
                // 批量处理
                print!("\n请输入图片目录: ");
                io::stdout().flush().unwrap();
                let input_dir = get_input();

                if input_dir.is_empty() {
                    println!("{} 目录不能为空", "[错误]".red().bold());
                    continue;
                }

                println!("\n{} 扫描图片...", "[扫描]".bright_cyan());
                let images = scan_images(&input_dir);

                if images.is_empty() {
                    println!("{} 未找到图片文件", "[警告]".yellow());
                    continue;
                }

                println!(
                    "{} 找到 {} 张图片",
                    "[成功]".green(),
                    images.len().to_string().bright_white().bold()
                );

                for (idx, img_path) in images.iter().enumerate() {
                    println!(
                        "  {}. {}",
                        idx + 1,
                        img_path.display().to_string().bright_black()
                    );
                }

                print!("\n输出目录 (默认: output): ");
                io::stdout().flush().unwrap();
                let output_dir = get_input();
                let output_dir = if output_dir.is_empty() {
                    "output"
                } else {
                    &output_dir
                };

                print!("\n确认处理 {} 张图片？(y/n): ", images.len());
                io::stdout().flush().unwrap();
                if get_input().to_lowercase() != "y" {
                    println!("已取消");
                    continue;
                }

                let processor = ImageProcessor::new(config.clone());
                println!("\n{}", "[开始] 开始批量处理...".bright_cyan().bold());

                match processor.process_batch(&images, output_dir) {
                    Ok(processed) => {
                        println!("\n{} 批量处理完成！", "[完成]".bright_green().bold());
                        println!(
                            "  {} 成功: {} 张",
                            "[成功]".green(),
                            processed.len().to_string().bright_white().bold()
                        );
                        println!("  输出目录: {}", output_dir.bright_cyan());
                    }
                    Err(e) => {
                        println!("{} 处理失败: {}", "[失败]".red().bold(), e);
                    }
                }
            }
            "3" => {
                configure_watermark(&mut config);
            }
            "4" => {
                display_config(&config);
            }
            "5" => {
                if let Err(e) = create_sample_image() {
                    println!("{} 创建失败: {}", "[失败]".red().bold(), e);
                }
            }
            "0" => {
                println!("\n{}", "感谢使用，再见！".bright_green().bold());
                break;
            }
            _ => println!("\n{} 无效选择，请重试", "[错误]".red().bold()),
        }
    }
}
